﻿using System;
using System.Diagnostics;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;
using System.Threading;
using System.Xml;
using Microsoft.VisualStudio.TestTools.Common;
using Microsoft.VisualStudio.TestTools.Execution;
using Microsoft.VisualStudio.TestTools.UITesting;

namespace PSRDataCollector
{
    // Problem Screen Recorder diagnostic data adapter
    [DataCollectorTypeUri("datacollector://Microsoft/PSRDataCollector/1.0")]
    [DataCollectorFriendlyName("Problem Screen Recorder", false)]
    public class PSRDataCollector : DataCollector
    {
        #region Constants

        private const string PSRExe = "psr.exe";
        private const int SaveDelay = 5000;

        #endregion Constants

        #region Fields

        private DataCollectionEvents dataEvents;
        private DataCollectionLogger dataLogger;
        private DataCollectionSink dataSink;
        private XmlElement configurationSettings;
        private string PSRLogFile;

        #endregion Fields

        #region DataCollector

        // Required method called by the testing framework
        public override void Initialize(XmlElement configurationElement,
                                        DataCollectionEvents events,
                                        DataCollectionSink sink,
                                        DataCollectionLogger logger,
                                        DataCollectionEnvironmentContext environmentContext)
        {
            dataEvents = events; // The test events
            dataLogger = logger; // The error and warning log
            dataSink = sink;     // Saves collected data
            // Configuration from the test settings
            configurationSettings = configurationElement;

            // Register common events for the data collector
            // Not all of the events are used in this class
            dataEvents.SessionStart +=
                new EventHandler<SessionStartEventArgs>(OnSessionStart);
            dataEvents.SessionEnd +=
                new EventHandler<SessionEndEventArgs>(OnSessionEnd);
            dataEvents.TestCaseStart +=
                new EventHandler<TestCaseStartEventArgs>(OnTestCaseStart);
            dataEvents.TestCaseEnd +=
                new EventHandler<TestCaseEndEventArgs>(OnTestCaseEnd);
            dataEvents.DataRequest +=
                new EventHandler<DataRequestEventArgs>(OnDataRequest);
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                dataEvents.SessionStart -=
                    new EventHandler<SessionStartEventArgs>(OnSessionStart);
                dataEvents.SessionEnd -=
                    new EventHandler<SessionEndEventArgs>(OnSessionEnd);
                dataEvents.TestCaseStart -=
                    new EventHandler<TestCaseStartEventArgs>(OnTestCaseStart);
                dataEvents.TestCaseEnd -=
                    new EventHandler<TestCaseEndEventArgs>(OnTestCaseEnd);
                dataEvents.DataRequest -=
                    new EventHandler<DataRequestEventArgs>(OnDataRequest);
            }
        }

        #endregion DataCollector

        #region Event Handlers

        public void OnSessionStart(object sender, SessionStartEventArgs e)
        {
            // TODO: Provide implementation
        }

        public void OnSessionEnd(object sender, SessionEndEventArgs e)
        {
            // TODO: Provide implementation
        }

        public void OnTestCaseStart(object sender, TestCaseEventArgs e)
        {
            this.PSRLogFile = Path.Combine(Environment.CurrentDirectory, e.TestCaseName + ".psr.zip");

            // Kill all running PSR processes
            TryKillProcess(PSRExe);

            // Start PSR
            try
            {
                InvokeProcess(PSRExe, String.Format("/start /output \"{0}\" /gui 0 /sc 1 /sketch 1 /maxsc 100", this.PSRLogFile));
            }
            catch (Exception exp)
            {
                dataLogger.LogError(e.Context, string.Format("Unexpected exception while trying to start PSR process : {0}", exp));
            }
        }

        public void OnTestCaseEnd(object sender, TestCaseEndEventArgs e)
        {
            try
            {
                // Save the PSR logs
                InvokeProcess(PSRExe, @"/stop").WaitForExit(60000);

                // Sleep to ensure PSR completes file creation operation
                Thread.Sleep(SaveDelay);

                if (!File.Exists(this.PSRLogFile))
                {
                    dataLogger.LogError(e.Context, "No user actions were recorded by PSR!");
                }
                else if (e.TestOutcome != TestOutcome.Passed)
                {
                    string ScreenshotFileName = Path.Combine(Environment.CurrentDirectory, e.TestCaseName + ".screenshot.jpg");
                    CaptureScreenshot(ScreenshotFileName);
                    dataSink.SendFileAsync(e.Context, ScreenshotFileName, false);

                    dataSink.SendFileAsync(e.Context, this.PSRLogFile, false);
                }
                else
                {
                    File.Delete(this.PSRLogFile);
                }
            }
            catch (Exception exp)
            {
                dataLogger.LogError(e.Context, string.Format("Unexpected exception while trying to stop PSR process : {0}", exp));
                // Kill all PSR processes
                TryKillProcess(PSRExe);
            }
        }

        public void OnDataRequest(object sender, DataRequestEventArgs e)
        {
            // TODO: Provide implementation
            // Most likely this occurs because a bug is being filed
        }

        #endregion Event Handlers

        #region Helpers

        private Process InvokeProcess(string processName, string parameters)
        {
            ProcessStartInfo startInfo = new ProcessStartInfo(processName, parameters);
            startInfo.WindowStyle = ProcessWindowStyle.Hidden;
            startInfo.UseShellExecute = true;
            startInfo.ErrorDialog = false;

            Process proc = new Process();
            proc.StartInfo = startInfo;
            proc.Start();

            return proc;
        }

        private void TryKillProcess(string processName)
        {
            Process[] processes = Process.GetProcessesByName(processName);
            foreach (Process proc in processes)
            {
                try { proc.Kill(); }
                catch (Exception exp) { }
            }
        }

        private void CaptureScreenshot(string ScreenshotFileName)
        {
            Playback.Initialize();
            Image img = UITestControl.Desktop.CaptureImage();
            Playback.Cleanup();
            img.Save(ScreenshotFileName, ImageFormat.Jpeg);
        }

        #endregion Helpers
    }
}